import astropy.io.fits as fits
import math
import numpy as np
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import astropy.units as u
import pickle
from scipy import interpolate
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)
c=3.0e10

def rms(x, axis=None):
    return np.sqrt(np.mean(x**2, axis=axis))


spectral_dict=pickle.load(open("spectral_dict.pickle", "rb"))

v_src=4.25
methanol_trans={'line1' : {'freq' : 241806.524e6, 'A_ul': 10**-4.66150,'E_u': 115.16019, 'g_u': 11, 'trans': '5(4)- - 4(4)-'},
                'line2' : {'freq' : 241813.255e6, 'A_ul': 10**-4.66204,'E_u': 122.72124, 'g_u': 11, 'trans': '5(-4) - 4(-4) E2'},
                'line3' : {'freq' : 241852.299e6, 'A_ul': 10**-4.40954,'E_u': 97.53028, 'g_u': 11 , 'trans': '5(-3) - 4(-3) E2'},
                'line4' : {'freq' : 241879.025e6, 'A_ul': 10**-4.22470,'E_u': 55.87058, 'g_u': 11, 'trans': '5(1) - 4(1) E1 '},
                'line5' : {'freq' : 241887.674e6, 'A_ul': 10**-4.29086,'E_u': 72.53337, 'g_u':11 , 'trans': '5(2)+ - 4(2)+ '}
               }

for key in spectral_dict.keys():
   fig,ax=plt.subplots(1,1,figsize=(15, 6))
   ax.plot(spectral_dict['CH3OH_241']['freq_axis']/1e9,spectral_dict['CH3OH_241']['spectrum'],color='black',drawstyle='steps',linewidth='1')
   ax.xaxis.set_major_locator(MultipleLocator(0.03))
   ax.xaxis.set_minor_locator(AutoMinorLocator())
   ax.tick_params(which='both', width=2)
   ax.tick_params(which='major', length=4)
   ax.tick_params(which='minor', length=2)
   ax.set_title('CH$_3$OH 241 GHz')
   ax.set_ylabel('Flux Density (Jy)')
   ax.set_xlabel('Frequency (GHz)',labelpad=-4)
   for key in methanol_trans.keys():
      dnu=v_src/3.0e5*methanol_trans[key]['freq']
      #ax.text((methanol_trans[key]['freq']-dnu)/1.0e9,0.25,'|')
      ax.plot([(methanol_trans[key]['freq']-dnu)/1.0e9,(methanol_trans[key]['freq']-dnu)/1.0e9],[0.25,0.27],color='black')
   ax.set_xlim(np.min(spectral_dict['CH3OH_241']['freq_axis']/1e9),np.max(spectral_dict['CH3OH_241']['freq_axis']/1e9))
   plt.savefig('CH3OH_241_spectrum.png')
   plt.savefig('CH3OH_241_spectrum.pdf')

#extract sub spectra
for key in methanol_trans.keys():
   dnu=v_src/3.0e5*methanol_trans[key]['freq']
   rel_freq=spectral_dict['CH3OH_241']['freq_axis']-methanol_trans[key]['freq']+dnu
   methanol_trans[key]['rel_freq_axis']=spectral_dict['CH3OH_241']['freq_axis']-methanol_trans[key]['freq']
   index=np.where(np.abs(rel_freq) < 4.5e6)
   sub_rel_freq=rel_freq[index[0]]
   sub_spectrum=spectral_dict['CH3OH_241']['spectrum'][index[0]]
   sub_e_spectrum=spectral_dict['CH3OH_241']['e_spectrum'][index[0]]
   sub_spectrum_jy_beam=spectral_dict['CH3OH_241']['spectrum_Jy_beam'][index[0]]
   sub_e_spectrum_jy_beam=spectral_dict['CH3OH_241']['e_spectrum_Jy_beam'][index[0]]
   methanol_trans[key]['sub_rel_freq']=sub_rel_freq.copy()
   methanol_trans[key]['sub_spectrum']=sub_spectrum.copy()
   methanol_trans[key]['sub_e_spectrum']=sub_e_spectrum.copy()
   methanol_trans[key]['sub_spectrum_Jy_beam']=sub_spectrum_jy_beam.copy()
   methanol_trans[key]['sub_e_spectrum_Jy_beam']=sub_e_spectrum_jy_beam.copy()
   methanol_trans[key]['nbeams']=spectral_dict['CH3OH_241']['nbeams']+0.0
   methanol_trans[key]['npix_per_beam']=spectral_dict['CH3OH_241']['npix_per_beam']+0.0
   index1=np.where(sub_rel_freq<-3.0e6)
   index2=np.where(sub_rel_freq>3.0e6)

   zero_low=np.mean(sub_spectrum[index1[0]])
   zero_high=np.mean(sub_spectrum[index2[0]])

   slope=(zero_low-zero_high)/(-3.5e6-3.5e6)
   y_intercept=(zero_low+zero_high)/2.0
   new_sub_spectrum=sub_spectrum-(slope*sub_rel_freq+y_intercept)

   zero_low=np.mean(sub_spectrum_jy_beam[index1[0]])
   zero_high=np.mean(sub_spectrum_jy_beam[index2[0]])

   slope=(zero_low-zero_high)/(-3.5e6-3.5e6)
   y_intercept=(zero_low+zero_high)/2.0
   new_sub_spectrum_jy_beam=sub_spectrum_jy_beam-(slope*sub_rel_freq+y_intercept)




   fig,ax=plt.subplots(1,1,figsize=(15, 6))
   ax.plot(sub_rel_freq/1e9,sub_spectrum,drawstyle='steps',linewidth='1')
   plt.savefig(key+'_spectrum.png')

   fig,ax=plt.subplots(1,1,figsize=(15, 6))
   ax.plot(sub_rel_freq/1e9,new_sub_spectrum,drawstyle='steps',linewidth='1')
   plt.savefig(key+'_new_spectrum.png')
   dnu=sub_rel_freq[1]-sub_rel_freq[0]
   dv=dnu/methanol_trans[key]['freq']*c
   methanol_trans[key]['dnu']=dnu.copy()
   methanol_trans[key]['dv']=dv.copy()
   methanol_trans[key]['integrated']=np.sum(new_sub_spectrum)
   methanol_trans[key]['integrated_Jy_beam']=np.sum(new_sub_spectrum_jy_beam)
   print(methanol_trans[key]['freq'],methanol_trans[key]['integrated'],methanol_trans[key]['dnu'],methanol_trans[key]['dv'])


pickle.dump(methanol_trans, open("methanol_transitions.pickle", "wb"))

#plot all sub spectra
fig,ax=plt.subplots(1,1,figsize=(15, 6))
for key in methanol_trans.keys():
   ax.plot(methanol_trans[key]['sub_rel_freq']/1e9,methanol_trans[key]['sub_spectrum'],drawstyle='steps',linewidth='1')
plt.savefig('ch3oh_sub_spectra.png')

#plot all sub spectra nomalized
fig,ax=plt.subplots(1,1,figsize=(15, 6))
for key in methanol_trans.keys():
   ax.plot(methanol_trans[key]['sub_rel_freq']/1e9,methanol_trans[key]['sub_spectrum']/np.max(methanol_trans[key]['sub_spectrum']),drawstyle='steps',linewidth='1')
plt.savefig('norm_ch3oh_sub_spectra.png')

#upsample spectrum and average
new_rel_freq=np.linspace(-4.0e6,4.0e6,401)
new_avg_spectrum=np.zeros(401)

fig,ax=plt.subplots(1,1,figsize=(15, 6))
for key in methanol_trans.keys():
   f=interpolate.interp1d(methanol_trans[key]['sub_rel_freq'],methanol_trans[key]['sub_spectrum'],kind='cubic')
   new_sub_spectrum=f(new_rel_freq)
   #new_sub_spectrum=np.interp(new_rel_freq,methanol_trans[key]['sub_rel_freq'],methanol_trans[key]['sub_spectrum'])
   ax.plot(new_rel_freq/1e9,new_sub_spectrum/np.max(new_sub_spectrum),drawstyle='steps',linewidth='1')
   new_avg_spectrum+=new_sub_spectrum/np.max(new_sub_spectrum)

plt.savefig('upsampled_ch3oh_sub_spectra.png')


fig,ax=plt.subplots(1,1,figsize=(15, 6))
ax.plot(new_rel_freq/1e9,new_avg_spectrum,drawstyle='steps',linewidth='1')
plt.savefig('upsampled_averag_ch3oh_sub_spectra.png')

index1=np.where(new_rel_freq<-3.0e6)
index2=np.where(new_rel_freq>3.0e6)

zero_low=np.mean(new_avg_spectrum[index1[0]])
zero_high=np.mean(new_avg_spectrum[index2[0]])

slope=(zero_low-zero_high)/(-3.5e6-3.5e6)
y_intercept=(zero_low+zero_high)/2.0
new_avg_spectrum=new_avg_spectrum-(slope*new_rel_freq+y_intercept)

new_avg_spectrum_wflip=(new_avg_spectrum+np.flip(new_avg_spectrum))/2.0

new_avg_spectrum_wflip=new_avg_spectrum_wflip/np.max(new_avg_spectrum_wflip)

fig,ax=plt.subplots(1,1,figsize=(15, 6))
ax.plot(new_rel_freq/1e9,new_avg_spectrum,drawstyle='steps',linewidth='1')
ax.plot(new_rel_freq/1e9,np.flip(new_avg_spectrum),drawstyle='steps',linewidth='1')
ax.plot(new_rel_freq/1e9,new_avg_spectrum_wflip,drawstyle='steps',linewidth='1')
plt.savefig('baseline_slope_upsampled_averag_ch3oh_sub_spectra.png')


spectral_template={'freq_axis': new_rel_freq.copy(),'spectrum': new_avg_spectrum_wflip.copy() }


pickle.dump(spectral_template, open("spectral_template.pickle", "wb"))









